#!/usr/bin/env python3

import sys
from elftools.elf.elffile import ELFFile
from elftools.dwarf.die import DIE, AttributeValue
from elftools.dwarf.compileunit import CompileUnit
# from tree_printer import print_list_tree
from typing import Optional

from elftools.dwarf.descriptions import ExprDumper, set_global_machine_arch

g_visitor: Optional[ExprDumper] = None

# def get_die_in_cu_by_offset(cu: CompileUnit, offset) -> Optional[DIE]:
#     for die in cu.iter_DIEs():  # type: DIE
#         if die.offset == cu.cu_offset + offset:
#             return die
#     return None


def get_die_in_cu_by_offset(cu: CompileUnit, offset) -> Optional[DIE]:
    left = 0
    right = len(cu._dielist) - 1
    real_offset = cu.cu_offset + offset

    while left <= right:
        mid = left + (right - left) // 2
        mid_entry: DIE = cu._get_DIE(mid)
        if mid_entry.offset < real_offset:
            left = mid + 1
        elif mid_entry.offset > real_offset:
            right = mid - 1
        else:
            return mid_entry
    return None


def get_location(entry: DIE) -> str:
    global g_visitor
    assert isinstance(g_visitor, ExprDumper)
    g_visitor.clear()

    location: AttributeValue = entry.attributes.get('DW_AT_location')
    if location is None:
        return '<No Loc Attr>'

    try:
        g_visitor.process_expr(location.value)
    except:
        return '<Visitor Failed>'

    return g_visitor.get_str()


def get_type_str(entry: DIE, cu: CompileUnit) -> str:
    return base_type_to_str(entry, cu)


def base_type_to_str(entry: DIE, cu: CompileUnit) -> str:
    if entry is None:
        return '{NoneDIE}'

    tag: str = entry.tag
    # if tag in ['DW_TAG_structure_type', 'DW_TAG_class_type', 'DW_TAG_base_type']:
    #     return entry.attributes.get('name', '{NoName}')

    base_type: AttributeValue = entry.attributes.get('DW_AT_type')
    if base_type is None:
        name_attr: AttributeValue = entry.attributes.get('DW_AT_name')
        if name_attr is None:
            return '{NoName}'
        else:
            return name_attr.value.decode('utf-8')

    base_entry = get_die_in_cu_by_offset(cu, base_type.value)

    fmt_str_dict = {
        'DW_TAG_const_type':            'const({})',
        'DW_TAG_typedef':               '{}',
        'DW_TAG_variable':              '{}',
        'DW_TAG_rvalue_reference_type': '{}',
        'DW_TAG_formal_parameter':      '{}',
        'DW_TAG_volatile_type':         '{}',
        'DW_TAG_array_type':            '{}[]',
        'DW_TAG_pointer_type':          '{}*',
        'DW_TAG_reference_type':        '{}&',
    }
    if tag in fmt_str_dict:
        return fmt_str_dict[tag].format(base_type_to_str(base_entry, cu))
    if tag == 'DW_TAG_enumeration_type':
        return 'int'

    return 'zzz ' + tag


def handle_file(file_path: str):
    print(f'Processing file: {file_path}')
    with open(file_path, 'rb') as f:
        elf_file = ELFFile(f)

        if not elf_file.has_dwarf_info():
            print('  file has no DWARF info')
            return

        # get_dwarf_info returns a DWARFInfo context object, which is the
        # starting point for all DWARF-based processing in pyelftools.
        dwarf_info = elf_file.get_dwarf_info()
        global g_visitor
        g_visitor = ExprDumper(dwarf_info.structs)
        set_global_machine_arch(elf_file.get_machine_arch())

        for CU in dwarf_info.iter_CUs():
            print('<', end='')
            for cu in dwarf_info.iter_CUs():
                print('0x%x' % cu.cu_offset, end=' ')
            print('>')
            print('Found a CU at offset 0x%x, die_offset 0x%x, length 0x%x %s' %
                  (CU.cu_offset, CU.cu_die_offset, CU['unit_length'], '*' * 80))

            for the_DIE in CU.iter_DIEs():  # type: DIE
                tag: str = the_DIE.tag
                if not isinstance(tag, str):
                    continue

                location_str = get_location(the_DIE)
                if location_str in ['<No Loc Attr>', '<Visitor Failed>']:
                    continue

                if tag in ['DW_TAG_variable', 'DW_TAG_formal_parameter']:
                    print('\n'.join([
                        f'',
                        'Offset: 0x%x' % the_DIE.offset,
                        f'TAG: {tag}\t\'{the_DIE.get_full_path()}\'',
                        f'Type: {get_type_str(the_DIE, CU)}',
                        f'Location: {location_str}',
                        f'Attrs: {the_DIE}',
                    ]))
                #else:
                #    print('\n'.join([
                #        f'',
                #        'Offset: 0x%x' % the_DIE.offset,
                #        f'TAG: {tag}\t\'{the_DIE.get_full_path()}\'',
                #        f'Attrs: {the_DIE}',
                #    ]))


if __name__ == '__main__':
    for filename in sys.argv[1:]:
        handle_file(filename)

